import os
import re
import datetime
import shutil
from datetime import datetime, timedelta

def get_files_from_folder(folder_path, extensions=['.mp4', '.mov', '.avi']):
    """获取文件夹中所有视频文件"""
    video_files = []
    for file in os.listdir(folder_path):
        if any(file.lower().endswith(ext) for ext in extensions):
            full_path = os.path.join(folder_path, file)
            mod_time = os.path.getmtime(full_path)
            video_files.append((full_path, mod_time))
    
    # 按修改时间排序
    video_files.sort(key=lambda x: x[1])
    return video_files

def parse_dji_filename(filename):
    """解析DJI文件名中的日期时间"""
    pattern = r'DJI_(\d{4})(\d{2})(\d{2})(\d{2})(\d{2})(\d{2})_'
    match = re.search(pattern, filename)
    
    if match:
        year = int(match.group(1))
        month = int(match.group(2))
        day = int(match.group(3))
        hour = int(match.group(4))
        minute = int(match.group(5))
        second = int(match.group(6))
        
        return datetime(year, month, day, hour, minute, second)
    return None

def generate_new_filename(old_filename, time_diff):
    """生成新的文件名，修正日期时间部分"""
    pattern = r'(DJI_)(\d{14})(_\d+_D.*)'
    match = re.search(pattern, old_filename)
    
    if match:
        prefix = match.group(1)
        date_time_str = match.group(2)
        suffix = match.group(3)
        
        # 解析日期时间字符串
        old_dt = datetime.strptime(date_time_str, '%Y%m%d%H%M%S')
        
        # 计算新的日期时间
        new_dt = old_dt + time_diff
        
        # 生成新的文件名
        new_date_time_str = new_dt.strftime('%Y%m%d%H%M%S')
        new_filename = f"{prefix}{new_date_time_str}{suffix}"
        
        return new_filename
    
    return old_filename  # 如果不匹配模式，返回原始文件名

def main():
    print("===== DJI相机时间修复工具 =====")
    folder_path = input("请输入视频文件所在的文件夹路径: ").strip()
    
    if not os.path.isdir(folder_path):
        print(f"错误: '{folder_path}' 不是有效的文件夹路径")
        return
    
    # 获取所有视频文件并按时间排序
    video_files = get_files_from_folder(folder_path)
    
    if not video_files:
        print(f"在 '{folder_path}' 中没有找到视频文件")
        return
    
    # 获取最早的文件
    earliest_file, earliest_time = video_files[0]
    earliest_filename = os.path.basename(earliest_file)
    
    print(f"\n找到最早的视频文件: {earliest_filename}")
    print(f"当前修改时间: {datetime.fromtimestamp(earliest_time)}")
    
    # 从文件名解析时间
    filename_time = parse_dji_filename(earliest_filename)
    if filename_time:
        print(f"从文件名解析的时间: {filename_time}")
    
    # 用户输入真实时间
    print("\n请输入该视频的真实拍摄时间:")
    try:
        real_year = int(input("年份 (例如 2023): "))
        real_month = int(input("月份 (1-12): "))
        real_day = int(input("日期 (1-31): "))
        real_hour = int(input("小时 (0-23): "))
        real_minute = int(input("分钟 (0-59): "))
        real_second = int(input("秒钟 (0-59): "))
        
        real_time = datetime(real_year, real_month, real_day, real_hour, real_minute, real_second)
    except ValueError as e:
        print(f"输入错误: {e}")
        return
    
    # 计算时间差
    if filename_time:
        # 使用文件名中的时间作为基准
        time_diff = real_time - filename_time
    else:
        # 使用文件修改时间作为基准
        time_diff = real_time - datetime.fromtimestamp(earliest_time)
    
    print(f"\n计算出的时间差: {time_diff}")
    
    # 确认是否继续
    confirm = input("\n是否确认修复所有文件? (y/n): ").strip().lower()
    if confirm != 'y':
        print("操作已取消")
        return
    
    # 创建备份文件夹
    backup_folder = os.path.join(folder_path, "backup_" + datetime.now().strftime("%Y%m%d%H%M%S"))
    os.makedirs(backup_folder, exist_ok=True)
    print(f"\n已创建备份文件夹: {backup_folder}")
    
    # 修复所有文件
    success_count = 0
    error_files = []
    
    for file_path, mod_time in video_files:
        try:
            filename = os.path.basename(file_path)
            
            # 备份原文件
            backup_path = os.path.join(backup_folder, filename)
            shutil.copy2(file_path, backup_path)
            
            # 计算新的修改时间
            new_mod_time = datetime.fromtimestamp(mod_time) + time_diff
            
            # 生成新的文件名
            new_filename = generate_new_filename(filename, time_diff)
            new_file_path = os.path.join(folder_path, new_filename)
            
            # 重命名文件
            if new_filename != filename:
                os.rename(file_path, new_file_path)
                file_path = new_file_path
            
            # 修改文件时间
            os.utime(file_path, (new_mod_time.timestamp(), new_mod_time.timestamp()))
            
            print(f"已修复: {filename} -> {new_filename} | 新时间: {new_mod_time}")
            success_count += 1
            
        except Exception as e:
            print(f"修复 {file_path} 时出错: {e}")
            error_files.append(file_path)
    
    print(f"\n操作完成! 成功修复 {success_count} 个文件")
    if error_files:
        print(f"有 {len(error_files)} 个文件修复失败")
    
    print(f"原始文件已备份至: {backup_folder}")

if __name__ == "__main__":
    main()
